#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动同步指定PR的修改，使用前需要安装GitPython
"""
import os
import sys
import json
import requests
import http.client
import git
from typing import List

########################################## 使用前注意修改以下配置 ######################################
LOCAL_REPO = "/home/zhaoyu/push_code/pytorch-npu"  # 你本地个人pta仓库路径
GITCODE_TOKEN = "XXXXXXXXXXXXXXXX"  # GitCode 私人令牌，需要从gitcode网站个人设置中获取
OFFICIAL_PR = "25488"   # 需要同步的Ascend/pytorch仓库里你自己创建的PR编号
BRANCHES = ["v2.6.0-7.2.0", "v2.7.1-7.2.0", "v2.8.0-7.2.0"] # 这些分支需要同步
UPSTREAM = "Ascend/pytorch" # PR会提交到这个目标仓库
MY_USER = "zhaoyu65"  # 你自己的gitcode用户名
MY_REPO = f"{MY_USER}/pytorch-npu"  # pytorch-npu是你自己的pta仓库名
#######################################################################################################


API = "https://gitcode.com/api/v4"
SESSION = requests.Session()
SESSION.headers.update({"PRIVATE-TOKEN": GITCODE_TOKEN})

def get_all_commit_id() -> list:
    conn = http.client.HTTPSConnection("api.gitcode.com")
    payload = ''
    headers = {
    'Accept': 'application/json'
    }
    conn.request("GET", f"/api/v5/repos/{UPSTREAM}/pulls/{OFFICIAL_PR}/commits?access_token={GITCODE_TOKEN}", payload, headers)
    res = conn.getresponse()
    data = res.read()
    text = data.decode("utf-8")
    # 把text的内容转换为json格式
    commit_id = []
    try:
        pr_data = json.loads(text)
        for pr in pr_data:
            commit_id.append(pr.get("sha"))

    except json.JSONDecodeError as e:
        print(f"❌ JSON解析失败: {e}")
        print(f"响应内容: {text}")
        return {}
    return commit_id

def get_pr_info() -> dict:
    # """获取 PR 的 commit-id / title / description"""
    conn = http.client.HTTPSConnection("api.gitcode.com")
    payload = ''
    headers = {
    'Accept': 'application/json'
    }
    conn.request("GET", f"/api/v5/repos/{UPSTREAM}/pulls/{OFFICIAL_PR}?access_token={GITCODE_TOKEN}", payload, headers)
    res = conn.getresponse()
    data = res.read()
    text = data.decode("utf-8")
    # 把text的内容转换为json格式
    try:
        pr_data = json.loads(text)
        return {
            "commit_id": get_all_commit_id(),
            "title": pr_data.get("title", ""),
            "description": pr_data.get("body", "")
        }
    except json.JSONDecodeError as e:
        print(f"❌ JSON解析失败: {e}")
        print(f"响应内容: {text}")
        return {}
    

def sync_one_branch(repo: git.Repo, branch: str, info: dict):
    """处理单个分支：拉最新代码 → cherry-pick → push → 创建 PR"""
    print(f"------ 处理 {branch} ------")
    # 1. 切到干净分支
    upstream_branch = f"upstream/{branch}"
    repo.git.checkout("HEAD", B=branch)
    repo.git.reset("--hard", upstream_branch)
    # 2. cherry-pick
    try:
        for id in info["commit_id"]:
            ret = repo.git.cherry_pick(id)
            print("repo.git.cherry_pick ret:", ret)
    except git.GitCommandError as e:
        print(f"⚠️  {branch} 冲突，需手动解决后执行：")
        print("   git cherry-pick --continue")
        print("   然后重新运行脚本")
        sys.exit(1)
    # 3. push
    origin = repo.remote("origin")
    # 使用带认证信息的push
    origin.push(branch, force=True)
    print(f"✅  {branch} 已推送")
    # 4. 创建 PR
    create_pr(branch, info)


def create_pr(source_branch: str, info: dict):
    """调用 GitCode API 创建 PR"""
    conn = http.client.HTTPSConnection("api.gitcode.com")
    payload = json.dumps({
        "title": info["title"],
        "head": f"{MY_USER}:{source_branch}",
        "base": source_branch,
        "body": info["description"],
        "milestone_number": 0,
        "labels": "",
        "issue": "",
        "assignees": "",
        "testers": "",
        "prune_source_branch": False,
        "draft": False,
        "squash": False,
        "squash_commit_message": "",
        "fork_path": MY_REPO
    })
    headers = {
    'Content-Type': 'application/json',
    'Accept': 'application/json'
    }
    conn.request("POST", f"/api/v5/repos/{UPSTREAM}/pulls?access_token={GITCODE_TOKEN}", payload, headers)
    res = conn.getresponse()
    data = res.read()
    text = data.decode("utf-8")
    # 把text的内容转换为json格式
    try:
        pr_data = json.loads(text)
        cid = pr_data.get("iid", "")
        if cid == "":
            print("❌ can not get commit id response")
        else:
            print("get commit id response:", cid)
            comment_pr(str(cid))
    except json.JSONDecodeError as e:
        print(f"❌ JSON解析失败: {e}")
        print(f"响应内容: {text}")


def comment_pr(commit_id: str):
    conn = http.client.HTTPSConnection("api.gitcode.com")
    payload = json.dumps({
    "body": "compile",
    "path": "",
    "position": 0
    })
    headers = {
    'Content-Type': 'application/json',
    'Accept': 'application/json'
    }
    conn.request("POST", f"/api/v5/repos/{UPSTREAM}/pulls/{commit_id}/comments?access_token={GITCODE_TOKEN}", payload, headers)
    res = conn.getresponse()
    data = res.read()
    print(data.decode("utf-8"))


def run_func():
    os.chdir(LOCAL_REPO)
    repo = git.Repo(".")
    # 在push之前设置凭据
    repo.git.config("credential.helper", "store")
    upstream_url = f"https://{MY_USER}:{GITCODE_TOKEN}@gitcode.com/{UPSTREAM}.git"
    origin_url   = f"https://{MY_USER}:{GITCODE_TOKEN}@gitcode.com/{MY_REPO}.git"
    # 确保远端存在
    if "upstream" not in [r.name for r in repo.remotes]:
        print(f"repo.create_remote")
        repo.create_remote("upstream", upstream_url)
    if "origin" not in [r.name for r in repo.remotes]:
        print("repo.create_remote")
        repo.create_remote("origin", origin_url)
    # 拉官方最新代码
    upstream = repo.remote("upstream")
    upstream.fetch()
    # 获取 PR 信息
    info = get_pr_info()
    # 循环处理每个分支
    for br in BRANCHES:
        sync_one_branch(repo, br, info)
    print("====== 全部完成 ======")

if __name__ == "__main__":
    run_func()